import scipy.io.wavfile as wav
from .config import *
from python_speech_features import logfbank, mfcc


# MFCC特征提取器
class MFCCExtractor:
    def __init__(self, sampling_rate=None, num_ceps=featureLen):
        """
        初始化 MFCCExtractor 类。

        参数:
        sampling_rate: int, 可选，默认为 None
            音频文件的采样率。如果为 None，则在读取音频文件时自动确定采样率。
        num_ceps: int, 默认为 13
            提取的 MFCC 系数的数量。
        """
        self.sampling_rate = sampling_rate
        self.num_ceps = num_ceps

    def extract_features(self, audio_path):
        """
        从给定的音频文件中提取 MFCC 特征。

        参数:
        audio_path: str
            音频文件的路径。

        返回:
        mfcc_feat: numpy.ndarray
            提取的 MFCC 特征。
        """
        # 读取音频文件
        if self.sampling_rate is None:
            rate, signal = wav.read(audio_path)
        else:
            rate = self.sampling_rate
            _, signal = wav.read(audio_path)

        # 提取 MFCC 特征
        mfcc_feat = mfcc(signal, rate, numcep=self.num_ceps, nfilt=featureLen)

        return mfcc_feat


# fbank特征提取器
class FbankExtractor:
    def __init__(self, sampling_rate=None):
        """
        初始化 FbankExtractor 类。

        参数:
        sampling_rate: int, 可选，默认为 None
            音频文件的采样率。如果为 None，则在读取音频文件时自动确定采样率。
        """
        self.sampling_rate = sampling_rate

    def extract_features(self, audio_path):
        """
        从给定的音频文件中提取 fbank 特征。

        参数:
        audio_path: str
            音频文件的路径。

        返回:
        fbank_feat: numpy.ndarray
            提取的 fbank 特征。
        """
        # 读取音频文件
        if self.sampling_rate is None:
            rate, signal = wav.read(audio_path)
        else:
            rate = self.sampling_rate
            _, signal = wav.read(audio_path)

        # 提取 fbank 特征
        fbank_feat = logfbank(signal, rate, nfilt=featureLen)

        return fbank_feat
